#ifndef _PREDICTOR_H_
#define _PREDICTOR_H_

#include "bezier_predict.h"
#include <traj_utils/planning_visualization.h>
#include <traj_utils/uniform_bspline.h>
#include <traj_utils/plan_container.hpp>
#include <unordered_map>

namespace ego_planner {
struct PredictedObj
{
  bool have_pos_{false};
  bool have_traj_{false};
  int max_size_{30};
  std::vector<Eigen::Vector4d> target_list_; // TODO: use deque, improve performance
  ObsPosPtr pos_{nullptr};                   // record the obs pos
  ObsTrajPtr traj_{nullptr};                 // record the obs traj

  int traj_order_{5};          // must be same as Bezierpredict::TRAJ_ORDER
  double t_start_gl_{0.0};     // global start time
  double prop_start_{0.0};     // propagation coefficient at time = poly_start_
  double R2_{0.0};             // adjust R2, goodness of fit
  Eigen::MatrixXd Evaluation_; // for calculate the variance of prediction

  PredictedObj() {}
  ~PredictedObj() {}
  PredictedObj(int max_size) : max_size_(max_size) {}

  void setInput(const Eigen::Vector4d &_target)
  {
    have_pos_ = true;

    target_list_.push_back(_target);

    if (target_list_.size() > max_size_)
      target_list_.erase(target_list_.begin());
  }

  // the same as Bezierpredict::getUncertainty
  double getUncertainty(double _t_now)
  {
    double t = _t_now - t_start_gl_;
    Eigen::VectorXd T = getT(t);
    double propagation_coef = T.dot(Evaluation_ * T) / prop_start_;
    return 1 - propagation_coef * (1 - R2_);
  }

  Eigen::VectorXd getT(double t)
  {
    Eigen::VectorXd T = Eigen::VectorXd::Zero(traj_order_ + 1);
    for (int i = 0; i < traj_order_ + 1; i++)
    {
      T(i) = pow(t, i);
    }
    return T;
  }
};

class Predictor
{

public:
  Predictor() {}
  ~Predictor() {}

  /* main API */
  void init(ros::NodeHandle &nh,
            Visualization::Ptr vis,
            ObsPosesPtr pos_ptr,
            ObsTrajsPtr traj_ptr);

  void setInput(const std::vector<int> &_ids,
                const std::vector<Eigen::Vector4d> &_targets);

  bool haveObs() { return have_obs_; }

  bool getPredictPos(std::vector<std::vector<Eigen::Vector3d>> &pos,
                     std::vector<std::vector<double>> &R2,
                     std::vector<std::vector<double>> &times,
                     std::vector<bool> &is_dangerous);

  typedef std::shared_ptr<Predictor> Ptr;
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW

private:
  typedef std::unordered_map<int, PredictedObj> ObjsMap;

  ObjsMap predicted_objs_;
  ObsPosesPtr obs_poses_{nullptr}; // record the obs pos
  ObsTrajsPtr obs_trajs_{nullptr}; // record the obs traj

  ros::Publisher vis_predict_list_pub_;
  ros::Publisher vis_control_pt_pub_;
  ros::Timer manager_timer_;

  bool have_obs_{false};
  size_t list_max_size_{30};
  size_t list_min_size_{7}; // must > 6

  Bezierpredict bezierPredict_;
  Visualization::Ptr visualization_;

  // paramter
  double sample_interval_{0.1};
  double time_remain_{1.0};   // if no target, remain 1s
  double time_vanish_{4.0};   // if no target, vanish after 4s
  double time_scale_{2.2};    // for predict
  double time_turncate_{4.0}; // turncate the time is too older

  void managerCb(const ros::TimerEvent &e);

  void updateObsPoses(int _id, const Eigen::Vector4d &_target);
  void updateObsTrajs(ObjsMap::iterator it, const ObsTraj &_obs_traj);
  void clearObsTraj(size_t _index);

  void updatePredict();
};

/* update the obs pos for other module */
inline void Predictor::updateObsPoses(
    int _id, const Eigen::Vector4d &_target)
{
  ObsPos _obs_pos;
  _obs_pos.id = _id;
  _obs_pos.time = ros::Time(_target(3));
  _obs_pos.pos = _target.head(3);

  ObjsMap::iterator it = predicted_objs_.find(_id);

  if (it == predicted_objs_.end()) {
    std::cerr << "can't find the index, BUG" << std::endl;
    return;
  }

  // before have no pos
  if (it->second.pos_ == nullptr)
  {
    it->second.pos_ = std::make_shared<ObsPos>(_obs_pos);
    obs_poses_->push_back(it->second.pos_);
  }
  else { // already have pos
    *(it->second.pos_) = _obs_pos;
  }
}

/* update the obs traj for other module */
inline void Predictor::updateObsTrajs(
    ObjsMap::iterator it, const ObsTraj &_obs_traj)
{
  it->second.have_traj_ = true;

  // before have no traj
  if (it->second.traj_ == nullptr)
  {
    it->second.traj_ = std::make_shared<ObsTraj>(_obs_traj);
    obs_trajs_->push_back(it->second.traj_);
  }
  else { // already have traj
    *(it->second.traj_) = _obs_traj;
  }
}

inline void Predictor::clearObsTraj(size_t _index)
{
  if (_index >= obs_trajs_->size()) return;

  int id = obs_trajs_->at(_index)->id;

  ObjsMap::iterator it = predicted_objs_.find(id);
  if (it != predicted_objs_.end())
  {
    it->second.have_traj_ = false;
    it->second.traj_ = nullptr;
  }
  else {
    std::cerr << "can't find the obs traj, BUG" << std::endl;
  }

  obs_trajs_->erase(obs_trajs_->begin() + _index);
}

} // namespace ego_planner
#endif